import pandas as pd

# 导入tushare
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

# 拉取数据
df = pro.daily(**{
    "ts_code": "600519.sh",
    "trade_date": "",
    "start_date": 20220414,
    "end_date": 20250414,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
print(df)

import pandas as pd
import numpy as np


# 日收益率
df['daily_return'] = df['close'].pct_change()

# 移动平均线
df['ma5'] = df['close'].rolling(window=5).mean()
df['ma10'] = df['close'].rolling(window=10).mean()
df['ma20'] = df['close'].rolling(window=20).mean()
df['ma60'] = df['close'].rolling(window=60).mean()

# 指数移动平均线
df['ema12'] = df['close'].ewm(span=12).mean()
df['ema26'] = df['close'].ewm(span=26).mean()

# MACD
df['macd'] = df['ema12'] - df['ema26']
df['signal'] = df['macd'].ewm(span=9).mean()
df['macd_hist'] = df['macd'] - df['signal']

# RSI
def calc_rsi(series, period=14):
    delta = series.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
    rs = gain / loss
    return 100 - (100 / (1 + rs))

df['rsi14'] = calc_rsi(df['close'])

# KDJ
low_min = df['low'].rolling(window=9).min()
high_max = df['high'].rolling(window=9).max()
df['rsv'] = (df['close'] - low_min) / (high_max - low_min) * 100
df['kdj_k'] = df['rsv'].ewm(com=2).mean()
df['kdj_d'] = df['kdj_k'].ewm(com=2).mean()
df['kdj_j'] = 3 * df['kdj_k'] - 2 * df['kdj_d']

# 布林带
df['boll_mid'] = df['close'].rolling(window=20).mean()
df['boll_std'] = df['close'].rolling(window=20).std()
df['boll_upper'] = df['boll_mid'] + 2 * df['boll_std']
df['boll_lower'] = df['boll_mid'] - 2 * df['boll_std']

# 波动率（20日标准差）
df['volatility20'] = df['daily_return'].rolling(window=20).std()

# 提取用于展示的技术指标列
df_show = df[['trade_date', 'close', 'ma5', 'ma10', 'macd', 'signal', 'rsi14', 'kdj_k', 'kdj_d', 'kdj_j', 'boll_upper', 'boll_lower', 'volatility20']]

# 只展示最近10条结果
print(df_show.tail(10))

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import numpy as np

# === 1. 构造标签（预测下一日收益率）===
df['next_return'] = df['daily_return'].shift(-1)

# === 2. 选取特征列（技术指标）===
features = ['ma5', 'ma10', 'macd', 'signal', 'rsi14', 'kdj_k', 'kdj_d', 'kdj_j', 'boll_upper', 'boll_lower', 'volatility20']
df_model = df.dropna(subset=features + ['next_return'])

X = df_model[features]
y = df_model['next_return']

# === 3. 拆分训练集和测试集（如70%训练）===
split = int(len(X) * 0.7)
X_train, X_test = X.iloc[:split], X.iloc[split:]
y_train, y_test = y.iloc[:split], y.iloc[split:]

# === 4. 随机森林模型训练 ===
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# === 5. 模型预测 ===
y_pred = model.predict(X_test)

# === 6. 评价回归模型性能 ===
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"MSE: {mse:.6f}")
print(f"R2 Score: {r2:.4f}")

# === 7. 构造简单交易策略 ===
# 如果预测收益率 > 0，则买入；否则不动
strategy_return = np.where(y_pred > 0, y_test.values, 0)

# 累计收益对比（策略 vs 实际买入并持有）
cumulative_strategy = np.cumprod(1 + strategy_return)
cumulative_market = np.cumprod(1 + y_test.values)

# === 8. 画图展示策略效果 ===
import matplotlib.pyplot as plt

plt.figure(figsize=(12, 6))
plt.plot(cumulative_market, label='Buy & Hold')
plt.plot(cumulative_strategy, label='Random Forest Strategy')
plt.legend()
plt.title('Cumulative Returns Comparison')
plt.xlabel('Days')
plt.ylabel('Cumulative Return')
plt.grid(True)
plt.tight_layout()
plt.show()

# 构建结果 DataFrame
result_df = pd.DataFrame({
'trade_date': df_model.iloc[split:]['trade_date'].values,
'actual_return': y_test.values,
'predicted_return': y_pred,
'strategy_signal': (y_pred > 0).astype(int),
'strategy_return': np.where(y_pred > 0, y_test.values, 0)
})

# 显示前10条结果
print(result_df.head(10))
